package org.ruoyi.util;

import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.boot.CommandLineRunner;
import org.springframework.core.io.FileSystemResource;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;

import java.io.File;
import java.util.HashSet;
import java.util.Set;


public class Dify implements CommandLineRunner {

    private static final String API_KEY = "dataset-w0wHDntaVokDZhB5FJnRhSK4"; // 填入你的 Key
    private static final String BASE_URL = "http://192.168.0.116:9000/v1";    // 填入你的 Base URL（注意 /v1 或 /api/v1）
    private static final String DATASET_NAME = "外科学1";
    private static final String DOCUMENT_FOLDER = "D:\\documents";
    // 允许上传的扩展名（小写）
    private static final String[] ALLOWED_EXTENSIONS = {".txt", ".pdf", ".docx"};
    // --- 配置结束 ---

    private final RestTemplate restTemplate = new RestTemplate();

    @Override
    public void run(String... args) {
        System.out.println("🚀 Dify 知识库自动上传任务开始");

        try {
            String datasetId = getOrCreateDatasetId();
            if (datasetId == null) {
                System.err.println("❌ 无法获取或创建知识库 ID，任务终止。");
                return;
            }

            File folder = new File(DOCUMENT_FOLDER);
            if (!folder.exists() || !folder.isDirectory()) {
                System.err.println("❌ 目标路径不是文件夹或不存在: " + DOCUMENT_FOLDER);
                return;
            }

            // 获取知识库中已有文档名，方便跳过
            Set<String> existingNames = listExistingDocumentNames(datasetId);

            System.out.println("📁 开始递归遍历目录并上传文件: " + DOCUMENT_FOLDER);
            traverseAndUpload(datasetId, folder, existingNames);

        } catch (Exception e) {
            System.err.println("⚠️ 任务执行异常: " + e.getMessage());
            e.printStackTrace();
        }

        System.out.println("✅ Dify 自动上传任务结束");
    }

    /**
     * 递归遍历目录并上传文件
     */
    private void traverseAndUpload(String datasetId, File dir, Set<String> existingNames) {
        File[] files = dir.listFiles();
        if (files == null) return;

        for (File f : files) {
            if (f.isDirectory()) {
                traverseAndUpload(datasetId, f, existingNames); // 递归子目录
                continue;
            }
            if (!isAllowedExtension(f.getName())) {
                System.out.println("跳过（扩展名不支持）: " + f.getAbsolutePath());
                continue;
            }

            String fileName = f.getName();
            if (existingNames.contains(fileName)) {
                System.out.println("跳过（已存在于知识库）: " + fileName);
                continue;
            }

            System.out.println("\n📤 上传文件: " + f.getAbsolutePath());
            boolean ok = uploadDocument(datasetId, f);
            if (ok) {
                existingNames.add(fileName); // 上传成功后加入集合，避免重复上传同名文件
            }
        }
    }

    private boolean isAllowedExtension(String name) {
        String lower = name.toLowerCase();
        for (String ext : ALLOWED_EXTENSIONS) {
            if (lower.endsWith(ext)) return true;
        }
        return false;
    }

    /**
     * 获取或创建知识库（返回 datasetId 或 null）
     */
    private String getOrCreateDatasetId() {
        String url = BASE_URL + "/datasets";
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + API_KEY);
        HttpEntity<Void> entity = new HttpEntity<>(headers);

        try {
            ResponseEntity<String> response = restTemplate.exchange(url, HttpMethod.GET, entity, String.class);
            System.out.println("获取知识库列表 状态码: " + response.getStatusCode());

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                JSONObject json = new JSONObject(response.getBody());
                JSONArray arr = json.optJSONArray("data");
                if (arr != null) {
                    for (int i = 0; i < arr.length(); i++) {
                        JSONObject ds = arr.getJSONObject(i);
                        if (DATASET_NAME.equals(ds.optString("name"))) {
                            String id = ds.optString("id");
                            System.out.println("✅ 知识库已存在: " + DATASET_NAME + " (ID: " + id + ")");
                            return id;
                        }
                    }
                } else {
                    // 有时候返回结构不一样，打印出来方便排查
                    System.out.println("注意：查询 datasets 返回结构不包含 data 字段，响应体: " + response.getBody());
                }
            } else {
                System.out.println("查询知识库列表失败，状态码: " + response.getStatusCode() + ", 响应: " + response.getBody());
            }

            System.out.println("📘 知识库不存在，开始创建...");
            return createDataset();

        } catch (Exception e) {
            System.err.println("检查或创建知识库时出错: " + e.getMessage());
            return null;
        }
    }

    /**
     * 创建知识库并返回 ID
     */
    private String createDataset() {
        String url = BASE_URL + "/datasets";
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + API_KEY);
        headers.setContentType(MediaType.APPLICATION_JSON);

        JSONObject body = new JSONObject();
        body.put("name", DATASET_NAME);
        body.put("permission", "all_team_members");

        HttpEntity<String> entity = new HttpEntity<>(body.toString(), headers);

        try {
            ResponseEntity<String> response = restTemplate.postForEntity(url, entity, String.class);
            System.out.println("创建知识库 状态码: " + response.getStatusCode());
            System.out.println("创建响应: " + response.getBody());

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                JSONObject json = new JSONObject(response.getBody());
                // Dify 的创建响应有时直接返回 { "id": "...", ... } 或 { "data": { "id": "..." } }
                if (json.has("id")) {
                    return json.getString("id");
                } else if (json.has("data") && json.getJSONObject("data").has("id")) {
                    return json.getJSONObject("data").getString("id");
                } else {
                    System.err.println("无法解析创建响应中的 id 字段: " + response.getBody());
                }
            } else {
                System.err.println("创建知识库失败，状态码: " + response.getStatusCode());
            }
        } catch (Exception e) {
            System.err.println("创建知识库失败: " + e.getMessage());
        }
        return null;
    }

    /**
     * 列出知识库中已存在的文档名称（用于跳过重复上传）
     */
    private Set<String> listExistingDocumentNames(String datasetId) {
        Set<String> names = new HashSet<>();
        try {
            String url = BASE_URL + "/datasets/" + datasetId + "/documents";
            HttpHeaders headers = new HttpHeaders();
            headers.set("Authorization", "Bearer " + API_KEY);
            HttpEntity<Void> entity = new HttpEntity<>(headers);

            ResponseEntity<String> response = restTemplate.exchange(url, HttpMethod.GET, entity, String.class);
            System.out.println("获取已上传文档列表 状态码: " + response.getStatusCode());

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                JSONObject json = new JSONObject(response.getBody());
                JSONArray arr = json.optJSONArray("data");
                if (arr != null) {
                    for (int i = 0; i < arr.length(); i++) {
                        JSONObject doc = arr.getJSONObject(i);
                        String name = doc.optString("name");
                        if (name != null && !name.isEmpty()) names.add(name);
                    }
                } else {
                    System.out.println("注意：documents 接口响应中没有 data 字段，响应体: " + response.getBody());
                }
            } else {
                System.err.println("获取文档列表失败，状态码: " + response.getStatusCode());
            }
        } catch (Exception e) {
            System.err.println("列出已上传文档时出错: " + e.getMessage());
        }
        return names;
    }

    /**
     * 上传单个文档，返回是否成功
     */
    private boolean uploadDocument(String datasetId, File fileToUpload) {
        String url = BASE_URL + "/datasets/" + datasetId + "/document/create-by-file";

        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + API_KEY);
        headers.setContentType(MediaType.MULTIPART_FORM_DATA);

        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        body.add("file", new FileSystemResource(fileToUpload));

        // data 字段（文本类型）
        JSONObject dataJson = new JSONObject();
        dataJson.put("indexing_technique", "high_quality");

        JSONObject processRule = new JSONObject();
        JSONObject rules = new JSONObject();

        JSONArray preProcessingRules = new JSONArray();
        preProcessingRules.put(new JSONObject().put("id", "remove_extra_spaces").put("enabled", true));
        preProcessingRules.put(new JSONObject().put("id", "remove_urls_emails").put("enabled", true));
        rules.put("pre_processing_rules", preProcessingRules);

        JSONObject segmentation = new JSONObject();
        segmentation.put("separator", "###");
        segmentation.put("max_tokens", 500);
        rules.put("segmentation", segmentation);

        processRule.put("rules", rules);
        processRule.put("mode", "custom");
        dataJson.put("process_rule", processRule);

        HttpHeaders dataHeaders = new HttpHeaders();
        dataHeaders.setContentType(MediaType.TEXT_PLAIN);
        HttpEntity<String> dataEntity = new HttpEntity<>(dataJson.toString(), dataHeaders);
        body.add("data", dataEntity);

        HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);

        try {
            ResponseEntity<String> response = restTemplate.postForEntity(url, requestEntity, String.class);
            System.out.println("上传响应 状态码: " + response.getStatusCode());
            System.out.println("响应体: " + response.getBody());
            if (response.getStatusCode().is2xxSuccessful()) {
                System.out.println("✅ 上传成功: " + fileToUpload.getName());
                return true;
            } else {
                System.err.println("❌ 上传失败 (" + fileToUpload.getName() + ")，状态码: " + response.getStatusCode());
                return false;
            }
        } catch (Exception e) {
            System.err.println("⚠️ 上传时发生 I/O 或网络错误 (" + fileToUpload.getName() + "): " + e.getMessage());
            return false;
        }
    }
}
